classdef MatDataLoader<handle
    %% public properties
    properties(Access=public)     
        batch_size        
        batch_index        
    end
    %% private properites
    properties(Access=private)
        dataset
        shuffle
        drop_last
        dataset_index
        dataset_len
    end
    %% public methods
    methods(Access=public)
        
        function obj = MatDataLoader(dataset, args)
            arguments
                dataset
                args.shuffle = true
                args.batch_size = 128
                args.drop_last = true
            end
            % parser input
            obj.dataset = dataset;
            obj.shuffle = args.shuffle;
            obj.batch_size = args.batch_size;
            obj.drop_last = args.drop_last;
            
            % data_len(int) - number of samples to be dealt with in the dataset;
            % If the last batch's size is less than batch_size, than drop
            % it.
            obj.dataset_len = obj.dataset.len();
            if obj.drop_last && mod(obj.dataset_len, obj.batch_size) ~= 0
                obj.dataset_len = obj.batch_size * floor(obj.dataset_len/obj.batch_size);
            end
            
            % dataset_index - the index of the samples to be read.
            obj.dataset_index = 1: obj.dataset.len;
            
            reset(obj)
            
        end
        
        function [x, y] = read(obj)
            if ~hasdata(obj)
                error('No data to read! Please reset the dataset.');
            end
            
            index = get_index(obj);
            [x, y] = obj.dataset.get_item(index);
            
            obj.batch_index = obj.batch_index+1;
           
        end
        
        function has_data = hasdata(obj)
             if obj.batch_index>length(obj)
                has_data = false;
             else 
                 has_data = true;
            end
        end
        
        function num_batch = length(obj)
            if obj.drop_last
                num_batch = floor(obj.dataset.len()/obj.batch_size);
            else
                num_batch = ceil(obj.dataset.len()/obj.batch_size);
            end
        end
        
         function reset(obj)
            obj.batch_index = 1;
            
            if obj.shuffle
                obj.dataset_index = randperm(obj.dataset.len(), obj.dataset_len);
            end
            
        end
        
    end
    %% private methods
    methods(Access=private)
        function index = get_index(obj)
            bs = obj.batch_size;
            bi = obj.batch_index;
            
            if bi == length(obj)
                index = obj.dataset_index((bi-1)*bs+1 :end);
            else
                index = obj.dataset_index((bi-1)*bs+1 :bi*bs);
            end
               
        end
              
    end
    
end
